package com.wolf.antlr;

import com.wolf.antlr.golang.GoLexer;
import com.wolf.antlr.golang.GoParser;
import com.wolf.antlr.golang.GoParserBaseListener;
import org.antlr.v4.runtime.CharStreams;
import org.antlr.v4.runtime.CommonTokenStream;
import org.antlr.v4.runtime.Lexer;
import org.antlr.v4.runtime.TokenStream;
import org.antlr.v4.runtime.tree.ParseTree;
import org.antlr.v4.runtime.tree.ParseTreeWalker;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class GolangSearcher {
    static List<MethodFragment> searchFunctions(String fileName, GolangListenerImp listener) {
        Lexer lexer;
        try {
            lexer = new GoLexer(CharStreams.fromFileName(fileName));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        TokenStream tokenStream = new CommonTokenStream(lexer);
        GoParser parser = new GoParser(tokenStream);
        ParseTree tree = parser.sourceFile();
        ParseTreeWalker walker = new ParseTreeWalker();
        walker.walk(listener, tree);
        return listener.getMethodFragments();
    }

    public static void main(String[] args) {
        GolangListenerImp listener = new GolangListenerImp();
        String fileName = Main.parserArgs(args, listener);
        List<MethodFragment> methodFragments = searchFunctions(fileName, listener);
        methodFragments.forEach(methodFragment -> System.out.printf("%s, %s, %s, %s, %s\n",
                methodFragment.startIndex(),
                methodFragment.endIndex(),
                methodFragment.startLine(),
                methodFragment.endLine(),
                methodFragment.methodDeclaration()));
    }

    static class GolangListenerImp extends GoParserBaseListener implements IListener {
        private String getMethodDeclaration(GoParser.MethodDeclContext ctx) {
            //获取函数名
            String func_name = ctx.IDENTIFIER().getText();
            // 获取接收者类型
            List<String> receiver_paras = new ArrayList<>();
            for (GoParser.ParameterDeclContext para : ctx.receiver().parameters().parameterDecl()) {
                receiver_paras.add(para.identifierList().getText() + " " + para.type_().getText());
            }

            List<String> signature_paras = new ArrayList<>();
            for (GoParser.ParameterDeclContext para : ctx.signature().parameters().parameterDecl()) {
                signature_paras.add(para.identifierList().getText() + " " + para.type_().getText());
            }
            // 获取返回类型
            String return_type = null;
            if (ctx.signature().result() != null) {
                return_type = ctx.signature().result().type_().getText();
            }
            // 拼接函数声明的完整格式
            StringBuilder func_decl = new StringBuilder();
            func_decl.append("func ");
            func_decl.append("(");
            func_decl.append(String.join(", ", receiver_paras));
            func_decl.append(") ");
            func_decl.append(func_name);
            func_decl.append("(");
            func_decl.append(String.join(", ", signature_paras));

            func_decl.append(")");
            if (return_type != null) {
                func_decl.append(" ");
                func_decl.append(return_type);
            }
            return func_decl.toString();
        }

        private String getMethodDeclaration(GoParser.FunctionDeclContext ctx) {
            String func_name = ctx.IDENTIFIER().getText();
            List<String> signature_paras = new ArrayList<>();
            for (GoParser.ParameterDeclContext para : ctx.signature().parameters().parameterDecl()) {
                signature_paras.add(para.identifierList().getText() + " " + para.type_().getText());
            }

            // 获取返回类型
            String return_type = null;
            if (ctx.signature().result() != null && ctx.signature().result().type_() != null) {
                return_type = ctx.signature().result().type_().getText();
            }
            // 拼接函数声明的完整格式
            StringBuilder func_decl = new StringBuilder();
            func_decl.append("func ");
            func_decl.append(func_name);
            func_decl.append("(");

            func_decl.append(String.join(", ", signature_paras));

            func_decl.append(")");

            if (return_type != null) {
                func_decl.append(" ");
                func_decl.append(return_type);
            }
            return func_decl.toString();
        }

        @Override
        public void enterMethodDecl(GoParser.MethodDeclContext ctx) {
            super.enterMethodDecl(ctx);
            String methodDeclaration = this.getMethodDeclaration(ctx);
            this.search(ctx, methodDeclaration);
        }

        @Override
        public void enterFunctionDecl(GoParser.FunctionDeclContext ctx) {
            super.enterFunctionDecl(ctx);
            String methodDeclaration = this.getMethodDeclaration(ctx);
            this.search(ctx, methodDeclaration);
        }
    }
}
